import time
from contextlib import asynccontextmanager
from contextvars import ContextVar
from typing import Callable, Iterator

from fastapi.exceptions import RequestValidationError, HTTPException
from fastapi.routing import APIRoute
from fastapi import Request, Response
import gzip


class GzipRequest(Request):
    async def body(self) -> bytes:
        if not hasattr(self, "_body"):
            body = await super().body()
            if "gzip" in self.headers.getlist("Content-Encoding"):
                body = gzip.decompress(body)
            self._body = body
        return self._body


class GzipRoute(APIRoute):
    def get_route_handler(self) -> Callable:
        original_route_handler = super().get_route_handler()

        async def custom_route_handler(request: Request) -> Response:
            request = GzipRequest(request.scope, request.receive)
            return await original_route_handler(request)

        return custom_route_handler


class ValidationErrorLoggingRoute(APIRoute):
    def get_route_handler(self) -> Callable:
        original_route_handler = super().get_route_handler()

        async def custom_route_handler(request: Request) -> Response:
            try:
                return await original_route_handler(request)
            except RequestValidationError as exc:
                body = await request.body()
                detail = {"infirmary_errors": exc.errors(), "body": body.decode()}
                raise HTTPException(status_code=422, detail=detail)

        return custom_route_handler


class TimedRoute(APIRoute):
    def get_route_handler(self) -> Callable:
        original_route_handler = super().get_route_handler()

        async def custom_route_handler(request: Request) -> Response:
            before = time.time()
            response: Response = await original_route_handler(request)
            duration = time.time() - before
            response.headers["X-Response-Time"] = str(duration)
            return response

        return custom_route_handler


class RequestRoute(APIRoute):
    # 指向当下的请求对象
    route_request: ContextVar[Request] = ContextVar("route_request")

    def get_route_request(self):
        self.route_request.get()

    @asynccontextmanager
    async def _set_route_request_token(self, request: Request) -> Iterator[None]:
        # 设置全局
        token = self.route_request.set(request)
        try:
            yield
        finally:
            self.route_request.reset(token)

    def get_route_handler(self) -> Callable:
        # 异步的方式调用
        original_route_handler = super().get_route_handler()

        async def custom_route_handler(request: Request) -> Response:
            async with self._set_route_request_token(request):
                response: Response = await original_route_handler(request)
                return response

        return custom_route_handler
